cfg = f"""
train:
  - /gemini/data-1/train

val:
  - /gemini/data-1/val

names:
  0: pig
  1: person
"""
with open("dataset.yml",'w') as fp:
    fp.write(cfg)


from ultralytics import YOLO

# Load a model
model = YOLO("yolo11s-seg.pt")  # load a pretrained model (recommended for training)

# Train the model
results = model.train(data="dataset.yml", save_dir="/gemini/output", epochs=100, imgsz=640, batch=0.9, cache=True)